from codesecurity.tasks.code_obfuscation_detect.cod import *
from codesecurity.data.api import *
from codesecurity.utils.speed_test import time_cost

import numpy as np
import time
import os
import fire
import shutil
import random


obfuscation_tags=['min','dist','lib','public','build']

def get_tag(path:str):   
    for tag in obfuscation_tags:
        if tag in path:
            return tag
    return 'other'

def get_package_ralated_path(target_dir,path:str):
    path=path[len(target_dir)+1:]
    start_index=path.find('/')
    path=path[start_index+1:]
    start_index=path.find('/')
    
    return path[start_index+1:]
    
@time_cost
def scan_obsfucation(target_dir,log_file,abstract_file):
    package_number=len(os.listdir(target_dir))

    files=iter_dir(target_dir)
    #files=[e for e in files if e.endswith('.js')]

    #print(f'find {len(files)} js files.')

    obs_count=0
    total=0
    list_e=[]

    for i,f in enumerate(files):
        total+=1
        e=ObfuscationDetector.entropy(f)
        if e:
            obs_count+=1
            list_e.append(f)
            if len(list_e)<1000:
                print(f)
                
        if i%10000==0:
            print(f'processed {i} files.')

    print(f'find {obs_count} obfuscated files.')


    with open(log_file,'w') as f:
        for e in list_e:
            f.write(e+'\n')
            
            
    statistic_dict={}
    prefix_dict={}

    for e in list_e:
        e_tag=get_tag(e)
        if e_tag not in statistic_dict:
            statistic_dict[e_tag]=0
        statistic_dict[e_tag]+=1
            
    for e in list_e:
        e=get_package_ralated_path(target_dir,e)
        prefix=os.path.dirname(e)
        if prefix not in prefix_dict:
            prefix_dict[prefix]=0
        prefix_dict[prefix]+=1

    with open(abstract_file,'w') as f:
        f.write(f'Searching {package_number} packages.\n')
        f.write(f'Find {obs_count} obfuscated files in {total} files.\n')
        f.write(f'Obfuscated files tags:\n')    
        for k,v in statistic_dict.items():
            f.write(f'{k}: {v}\n')
    
        f.write(f'Obfuscated files prefix:\n')
        for k,v in prefix_dict.items():
            f.write(f'{k}: {v}\n')


def collect_obfucation_file(list_file,target_dir,common_prefix=""):
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)
    
    with open(list_file,'r') as f:
        for line in f:
            line=line.strip()
            
            new_filename=line
            
            if len(common_prefix)>0:
                new_filename=new_filename[len(common_prefix):]
                
            new_filename=new_filename.replace('/','_')
            
            shutil.copy(line,os.path.join(target_dir,new_filename))
            print(f'copy {line} to {os.path.join(target_dir,new_filename)}')
            
def collect_normal_js_file(source_dir,target_dir):
    
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)
    
    all_files=iter_dir(source_dir)
    
    number=0
    
    
    print(f'try copy {source_dir}/[sample.js] to {os.path.join(target_dir,"[sample.js]")}')
    
    for line in all_files:

        if line.endswith('.js')==False:
            continue

        new_filename=line[len(source_dir):]
            
        new_filename=new_filename.replace('/','_')
        
        if len(new_filename)>255:
            new_filename=new_filename[len(new_filename)-255:]
        
        shutil.copy(line,os.path.join(target_dir,new_filename))
        
        if random.random()<0.01:
            print(f'processed {number} files.')
        
        number+=1
        
    print(f'collect finish. copy {number} files to {target_dir}.')
    
if __name__=='__main__':
    fire.Fire()

